import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import matplotlib
matplotlib.use('Agg')

from tqdm import tqdm
from configs.config import load_config
from data_process.tensor_io import TensorIO
from data_process.rtk_reader import RtkReader
from data_process.plot_bev import BEVGenerator
from data_process.annotation_processor import BEVAnnotationProcessor


if __name__ == "__main__":
    cfg = load_config("./configs/data_process.py")

    rtk_reader = RtkReader(cfg)

    json_root_dir = cfg.get('bev_output_folder')
    data_root_list = cfg.get('data_root_list')
    subdirs = [d for d in os.listdir(json_root_dir) if os.path.isdir(os.path.join(json_root_dir, d))]
    
    for subdir in tqdm(subdirs, desc="Processing folders", unit="folder"):
        subdir_path = os.path.join(json_root_dir, subdir)
        data_name = subdir.split('_')[0]
        cfg['data_name'] = subdir
        
        bev_json_path = os.path.join(json_root_dir, subdir, f'{subdir}.json')
        bev_meta_path = os.path.join(json_root_dir, subdir, f'total_BEV_{subdir}_meta.json')
        
        if not (os.path.exists(bev_json_path) and os.path.exists(bev_meta_path)):
            continue
        
        annotation, ref_rtk, meta = BEVAnnotationProcessor.load_annotation(
            bev_json_path,
            bev_meta_path
            )
        
        rtk_data = rtk_reader.read_tensor_rtk(os.path.join(cfg.get('data_root_list')[data_name], 'SARRadar'), subdir)

        label = BEVAnnotationProcessor.generate_tensor_labels(annotation, rtk_data, ref_rtk, cfg)
        TensorIO.save_labels(label, cfg)
        
    if cfg.get('debug', False):
        BEVGenerator.create_combine_bev(cfg)
            